#include <iostream>
#include <windows.h>

// 定义必要结构（因为 WinSDK 不自带）
typedef struct _PROCESS_BASIC_INFORMATION {
	PVOID Reserved1;
	PVOID PebBaseAddress;       // 🔥 就是 PEB 地址
	PVOID Reserved2[2];
	ULONG_PTR UniqueProcessId;
	PVOID Reserved3;
} PROCESS_BASIC_INFORMATION;

typedef enum _PROCESSINFOCLASS {
	ProcessBasicInformation = 0
} PROCESSINFOCLASS;

typedef NTSTATUS(WINAPI* pNtQueryInformationProcess)(
	HANDLE ProcessHandle,
	PROCESSINFOCLASS ProcessInformationClass,
	PVOID ProcessInformation,
	ULONG ProcessInformationLength,
	PULONG ReturnLength
	);

class shell {
	public:
		void* encoding(void* pePtr);
		void* decryption(void* pePtr);
		IMAGE_DOS_HEADER* getDosHead(void* pePtr);
		IMAGE_NT_HEADERS32* getNtHead(void* pePtr);
		DWORD getSizeOfImage(void* pePtr);
		DWORD getImageBase(void* pePtr);
		DWORD getSizeOfHeaders(void* pePtr);
		DWORD getNumberOfSections(void* pePtr);
		void* RvaToFileOffset(void* pePtr, DWORD rvaVlue);
		void* FoaToRvaOffset(void* pePtr, DWORD rvaValue);
		IMAGE_SECTION_HEADER* getSectionHeadStart(void* pePtr);
		void* addSection(void* pePtr, DWORD sizeValue);
		PVOID GetRemoteImageBase(HANDLE hProcess);
		
};

IMAGE_DOS_HEADER* shell::getDosHead(void* pePtr) {
	
	return (IMAGE_DOS_HEADER*)pePtr;
}

IMAGE_NT_HEADERS32* shell::getNtHead(void* pePtr) {
	IMAGE_DOS_HEADER* peDosPtr = (IMAGE_DOS_HEADER*)pePtr;
	IMAGE_NT_HEADERS32* nt_head_ptr = (IMAGE_NT_HEADERS32*)((DWORD)pePtr + peDosPtr->e_lfanew);
	return nt_head_ptr;
}

DWORD shell::getSizeOfImage(void* pePtr) {
	IMAGE_DOS_HEADER* peDosPtr = (IMAGE_DOS_HEADER*)pePtr;
	IMAGE_NT_HEADERS32* nt_head_ptr = (IMAGE_NT_HEADERS32*)((DWORD)pePtr + peDosPtr->e_lfanew);
	DWORD sizeOfImage = nt_head_ptr->OptionalHeader.SizeOfImage;
	return sizeOfImage;
}

DWORD shell::getImageBase(void* pePtr) {
	IMAGE_DOS_HEADER* peDosPtr = (IMAGE_DOS_HEADER*)pePtr;
	IMAGE_NT_HEADERS32* nt_head_ptr = (IMAGE_NT_HEADERS32*)((DWORD)pePtr + peDosPtr->e_lfanew);
	DWORD imageBase = nt_head_ptr->OptionalHeader.ImageBase;
	return imageBase;
}

DWORD shell::getSizeOfHeaders(void* pePtr) {
	IMAGE_DOS_HEADER* peDosPtr = (IMAGE_DOS_HEADER*)pePtr;
	IMAGE_NT_HEADERS32* nt_head_ptr = (IMAGE_NT_HEADERS32*)((DWORD)pePtr + peDosPtr->e_lfanew);
	DWORD sizeOfHeaders = nt_head_ptr->OptionalHeader.SizeOfHeaders;
	return sizeOfHeaders;
}

DWORD shell::getNumberOfSections(void* pePtr) {
	IMAGE_DOS_HEADER* peDosPtr = (IMAGE_DOS_HEADER*)pePtr;
	IMAGE_NT_HEADERS32* nt_head_ptr = (IMAGE_NT_HEADERS32*)((DWORD)pePtr + peDosPtr->e_lfanew);
	DWORD numberOfSections = nt_head_ptr->FileHeader.NumberOfSections;
	return numberOfSections;
}

IMAGE_SECTION_HEADER* shell::getSectionHeadStart(void* pePtr) {
	IMAGE_DOS_HEADER* dosPtr = getDosHead(pePtr);
	IMAGE_NT_HEADERS32* nt_head_ptr = getNtHead(pePtr);
	WORD sizeOfOptionalHeader = nt_head_ptr->FileHeader.SizeOfOptionalHeader;
	IMAGE_SECTION_HEADER* section_start_ptr = (IMAGE_SECTION_HEADER*)((DWORD)pePtr + dosPtr->e_lfanew + 4 + sizeof(IMAGE_FILE_HEADER) + sizeOfOptionalHeader);
	return section_start_ptr;
	
}

void* shell::RvaToFileOffset(void* pePtr, DWORD rvaVlue) {
	IMAGE_SECTION_HEADER* section_head = getSectionHeadStart(pePtr);
	DWORD sectionNumber = getNumberOfSections(pePtr);
	
	IMAGE_SECTION_HEADER* final_section_ptr = NULL;
	for (int i = 0; i < sectionNumber; i++) {
		IMAGE_SECTION_HEADER* current_section = section_head + i;
		DWORD section_start = current_section->VirtualAddress;
		DWORD section_size = current_section->Misc.VirtualSize > current_section->SizeOfRawData ? current_section->Misc.VirtualSize : current_section->SizeOfRawData;
		DWORD section_end = section_start + section_size;
		if (rvaVlue >= section_start && rvaVlue <= section_end) {
			final_section_ptr = current_section;
			break;
		}
	}

	if (final_section_ptr == NULL) {
		return NULL;
	}
	DWORD final_offset = final_section_ptr->PointerToRawData + (rvaVlue - final_section_ptr->VirtualAddress);
	return (void*)(final_offset);
}

void* shell::FoaToRvaOffset(void* pePtr, DWORD rvaValue) {
	IMAGE_SECTION_HEADER* section_head = getSectionHeadStart(pePtr);
	DWORD sectionNumber = getNumberOfSections(pePtr);

	IMAGE_SECTION_HEADER* final_section_ptr = NULL;
	for (int i = 0; i < sectionNumber; i++) {
		IMAGE_SECTION_HEADER* current_section = section_head + i;
		DWORD section_start = current_section->PointerToRawData;
		DWORD section_end = current_section->PointerToRawData + current_section->SizeOfRawData;
		if (rvaValue >= section_start && rvaValue <= section_end) {
			final_section_ptr = current_section;
			break;
		}
	}

	if (final_section_ptr == NULL) {
		return NULL;
	}
	DWORD final_offset = final_section_ptr->VirtualAddress + (rvaValue - final_section_ptr->VirtualAddress);
	return (void*)final_offset;
}

void* shell::addSection(void* pePtr, DWORD sizeValue) {
	IMAGE_NT_HEADERS32* nt_head = getNtHead(pePtr);
	IMAGE_SECTION_HEADER* section_head = getSectionHeadStart(pePtr);
	DWORD sizeOfImage = getSizeOfImage(pePtr);
	DWORD sectionNumber = getNumberOfSections(pePtr);
	IMAGE_SECTION_HEADER* lastSection = section_head + (sectionNumber - 1);
	IMAGE_SECTION_HEADER* newSection = section_head + sectionNumber;
	char* free_space_ptr = (char*)newSection;
	for (int i = 0; i < 80; i++)
	{
		if (*(free_space_ptr + i ) != 0) {
			printf("剩余空间不足！此空间有代码！！！");
			return;
			}
	}

	char name[] = ".tttt";
	char* sectionName = (char*)newSection->Name;
	strncpy(sectionName, name, 8);
	newSection->Misc.VirtualSize = sizeValue;
	newSection->SizeOfRawData = sizeValue;
	newSection->VirtualAddress = sizeOfImage;
	newSection->PointerToRawData = lastSection->PointerToRawData + lastSection->SizeOfRawData;
	newSection->Characteristics = 0x60000020;

	nt_head->FileHeader.NumberOfSections = nt_head->FileHeader.NumberOfSections + 1;
	nt_head->OptionalHeader.SizeOfImage = sizeOfImage + sizeValue;

	void* newPePtr = malloc(sizeOfImage + sizeValue);
	memcpy(newPePtr, pePtr, sizeOfImage + sizeValue);
	return newPePtr;
}

void* shell::encoding(void* pePtr) {
	
	// 先读取需要被加壳文件大小
	DWORD originSizeOfImage = getSizeOfImage(pePtr);
	// 加密,采用异或加密
	char* encodeStr = (char*)pePtr;
	for (int i = 0; i < originSizeOfImage; i++) {
		*(encodeStr + i) ^= 0x66;
	}
	
	// 这里需要读取壳子
	void* shellPePtr = NULL;
	// 新增节
	// 然后开辟一个新的大小节
	// 加密 把 这个内容加到最后一个节里
	void* newPePtr = addSection(shellPePtr, originSizeOfImage);
	IMAGE_SECTION_HEADER* newSectionStart = getSectionHeadStart(newPePtr);
	DWORD numberOfSection = getNumberOfSections(newPePtr);
	IMAGE_SECTION_HEADER* lastSectionStart = newSectionStart + (numberOfSection - 1);
	void* encode_start = (void*)((DWORD)newPePtr + lastSectionStart->PointerToRawData);
	memcpy(encode_start, encodeStr, originSizeOfImage);

	// 写到文件
	//writeExe()
}

// 获取远程ImageBase
PVOID shell::GetRemoteImageBase(HANDLE hProcess) {
	PROCESS_BASIC_INFORMATION pbi = { 0 };
	typedef NTSTATUS(WINAPI* pNtQueryInformationProcess)(
		HANDLE, PROCESSINFOCLASS, PVOID, ULONG, PULONG);
	auto NtQueryInformationProcess =
		(pNtQueryInformationProcess)GetProcAddress(
			GetModuleHandle(L"ntdll.dll"), "NtQueryInformationProcess");

	NtQueryInformationProcess(hProcess, ProcessBasicInformation, &pbi, sizeof(pbi), NULL);

	PVOID imageBase = NULL;
	ReadProcessMemory(hProcess, (BYTE*)pbi.PebBaseAddress + 0x10, &imageBase, sizeof(PVOID), NULL);
	return imageBase;
}

void* shell::decryption(void* pePtr) {
	
	
	// 读取这个壳程序最后的一个节
	// 解密这个最后一个节
	// 读取这个最后一个节，拉伸,并获取关键信息
	// 开启一个新进程，挂起
	// 获取线程上下文
	// 卸载外壳程序的文件镜像(ZwUnmapViewOfSection)
	// 在指定的位置(src的ImageBase)申请指定大小(src的SizeOfImage)的内存(VirtualAllocEx)
	// 如果创建失败，查看src是否包含重定位表，如果包含重定位表，就在任意位置申请(src的SizeOfImage)大小的内存，然后修复重定位表.
	// 如果在指定位置申请内存失败，并且没有重定位表的数据，直接返回失败.
	// 如果内存申请成功，将新的数据复制到内存中 
	// 修正运行环境的基址和入口地址 
	// 恢复运行

	IMAGE_SECTION_HEADER* sectionStart = getSectionHeadStart(pePtr);
	DWORD numberOfSection = getNumberOfSections(pePtr);
	IMAGE_SECTION_HEADER* lastSectionStart = sectionStart + (numberOfSection - 1);
	char* decodeSrcPtr = (char*)((DWORD)pePtr + lastSectionStart->PointerToRawData);
	DWORD encodeSrcSize = lastSectionStart->SizeOfRawData >= lastSectionStart->Misc.VirtualSize ? lastSectionStart->SizeOfRawData : lastSectionStart->Misc.VirtualSize;
	// 解密
	for (int i = 0; i < encodeSrcSize; i++) {
		*(decodeSrcPtr + i) ^= 0x66;
	}
	
	// 拉伸
	// 挂起方式启动一个EXE
	STARTUPINFO si = { 0 };
	PROCESS_INFORMATION pi;
	si.cb = sizeof(STARTUPINFO);
	TCHAR szBuffer[256] = TEXT("C://Program Files (x86)//Microsoft//Edge//Application//msedge.exe");
	CreateProcess(szBuffer, NULL, NULL, NULL, FALSE, CREATE_SUSPENDED, NULL, NULL, &si, &pi);
	CONTEXT contx;
	contx.ContextFlags = CONTEXT_FULL;
	GetThreadContext(pi.hThread, &contx);

	// 卸载
	typedef LONG(WINAPI* ZWUNMAPVIEWOFSECTION)(HANDLE, PVOID);
	ZWUNMAPVIEWOFSECTION ZwUnmapViewOfSection =
		(ZWUNMAPVIEWOFSECTION)GetProcAddress(GetModuleHandleW(L"ntdll.dll"), "ZwUnmapViewOfSection");
	PVOID baseAddress = GetRemoteImageBase(pi.hProcess);
	LONG status = ZwUnmapViewOfSection(pi.hProcess, baseAddress);
	if (status != 0) {
		std::cout << "ZwUnmapViewOfSection failed, status: " << std::hex << status << "\n";
		return NULL;
	}

	std::cout << "映像卸载成功，现在可以 VirtualAllocEx + WriteProcessMemory 写入你解密的PE 了。\n";

	DWORD imageBase = getImageBase((void*)decodeSrcPtr);

	//分配空间
	LPVOID allocAddress = VirtualAllocEx(
		pi.hProcess,            // pi 是 PROCESS_INFORMATION 中的进程句柄
		(LPVOID)imageBase,              // 建议分配的地址（可以是 imageBase 或 NULL）
		encodeSrcSize,                 // 映像大小（从原始 PE 文件读取）
		MEM_COMMIT | MEM_RESERVE,
		PAGE_EXECUTE_READWRITE  // 可读可写可执行
	);

	if (!allocAddress) {
		std::cout << "VirtualAllocEx 失败，错误码：" << GetLastError() << std::endl;
		// 判断是否有重定位表，如果有重定位表，就随机分配空间，没有重定位表直接报错
		//分配空间
		allocAddress = VirtualAllocEx(
			pi.hProcess,            // pi 是 PROCESS_INFORMATION 中的进程句柄
			NULL,              // 建议分配的地址（可以是 imageBase 或 NULL） NULL 是随机分配
			encodeSrcSize,                 // 映像大小（从原始 PE 文件读取）
			MEM_COMMIT | MEM_RESERVE,
			PAGE_EXECUTE_READWRITE  // 可读可写可执行
		);
	}
	else {
		std::cout << "在目标进程中分配内存成功，地址：" << allocAddress << std::endl;
	}



}

int main() {
	shell sg;
	void* ptr = (void*)GetModuleHandle(NULL);
	
	return 0;
}